//-----------------------------------------------------------------------------
// Author - Marco Martinelli - https://github.com/martinellimarco
//-----------------------------------------------------------------------------

#pragma warning (disable : 4127)

#include "scream.h"
#include "ivshmemsavedata.h"

#pragma code_seg("PAGE")
//=============================================================================
// CIVSHMEMSaveData
//=============================================================================

//=============================================================================
CIVSHMEMSaveData::CIVSHMEMSaveData() :
    m_ulChunkOffset(0),
    m_pStreamHdr(NULL),
    m_pSharedMem(NULL),
    m_SharedMemSize(0),
    m_fCapture(FALSE)
{
    PAGED_CODE();

    DPF_ENTER(("[CIVSHMEMSaveData::CIVSHMEMSaveData]"));
} // CIVSHMEMSaveData

//=============================================================================
CIVSHMEMSaveData::~CIVSHMEMSaveData() {
    PAGED_CODE();

    DPF_ENTER(("[CIVSHMEMSaveData::~CIVSHMEMSaveData]"));
    if (m_pSharedMem) {
        m_pStreamHdr->started = 0;
    }
} // CIVSHMEMSaveData

//=============================================================================
void CIVSHMEMSaveData::Disable(BOOL fDisable) {
    PAGED_CODE();

    if (m_pStreamHdr) {
        KeMemoryBarrier();
        m_pStreamHdr->started = !fDisable;
        m_pStreamHdr->fmt.fmtGeneration++;
        if (!fDisable) {
            m_ivshmem.chunkIdx = m_pStreamHdr->chunkIdx;
        }
    }
} // Disable

//=============================================================================
NTSTATUS CIVSHMEMSaveData::Initialize(PVOID pSharedMem, UINT uSharedMemSize,
                                      BOOLEAN fCapture, DWORD nSamplesPerSec,
                                      WORD wBitsPerSample, WORD nChannels,
                                      DWORD dwChannelMask)
{
    PAGED_CODE();

    NTSTATUS          ntStatus = STATUS_SUCCESS;

    DPF_ENTER(("[CIVSHMEMSaveData::Initialize]"));
    m_pSharedMem = pSharedMem;
    m_SharedMemSize = uSharedMemSize;
    if (!m_pSharedMem) {
        return STATUS_DEVICE_HARDWARE_ERROR;
    }

    m_fCapture = fCapture;

    m_ivshmem.chunkSize = (UINT32)((wBitsPerSample>>3)*nChannels*nSamplesPerSec/50);

    PIVSHMEM_SCREAM_HEADER shdr = (PIVSHMEM_SCREAM_HEADER)m_pSharedMem;
    if (fCapture) {
        m_pStreamHdr = &shdr->capt;
    } else {
        m_pStreamHdr = &shdr->play;
    }

    m_pStreamHdr->started = 0;
    m_pStreamHdr->chunkIdx = 0;
    m_ivshmem.chunkIdx = 0;
    m_ulChunkOffset = 0;

    // Half of the memory is reserved for play, another half for read.
    ULONG bufSize = (ULONG)(m_SharedMemSize - sizeof(IVSHMEM_SCREAM_HEADER)) / 2;

    m_pStreamHdr->chunkSize = m_ivshmem.chunkSize;
    m_pStreamHdr->maxChunks = m_ivshmem.maxChunks = (UINT16)(bufSize / m_pStreamHdr->chunkSize);

    m_ivshmem.bufferSize = bufSize;

    if (fCapture) {
        m_pStreamHdr->offset = m_ivshmem.offset = sizeof(IVSHMEM_SCREAM_HEADER) + bufSize;
    } else {
        m_pStreamHdr->offset = m_ivshmem.offset = sizeof(IVSHMEM_SCREAM_HEADER);
    }

    // Only multiples of 44100 and 48000 are supported
    m_pStreamHdr->fmt.sampleRate = (UINT8)((nSamplesPerSec % 44100) ? (0 + (nSamplesPerSec / 48000)) : (128 + (nSamplesPerSec / 44100)));
    m_pStreamHdr->fmt.sampleSize = (UINT8)(wBitsPerSample);
    m_pStreamHdr->fmt.channels = (UINT8)(nChannels);
    m_pStreamHdr->fmt.channelMap = dwChannelMask;
    m_pStreamHdr->dmaStartTime = 0;
    m_pStreamHdr->started = FALSE;
    shdr->magic = MAGIC_NUM;

    return ntStatus;
} // Initialize

#pragma code_seg()
//=============================================================================
void CIVSHMEMSaveData::WriteData(IN PBYTE pBuffer, IN ULONG ulByteCount) {
    ASSERT(pBuffer);

    if (!m_pStreamHdr) {
        return;
    }

    if (!m_pStreamHdr->started) {
        return;
    }

    //DPF_ENTER(("[CIVSHMEMSaveData::WriteData ulByteCount=%lu]", ulByteCount));

    // Undersized (paranoia)
    if (0 == ulByteCount) {
        return;
    }

    // Oversized (paranoia)
    if (ulByteCount > (m_ivshmem.bufferSize / 2)) {
        return;
    }

    PBYTE shmem_ptr = (PBYTE)m_pSharedMem + m_pStreamHdr->offset;
    ULONG toWrite = ulByteCount;
    PBYTE pReadPtr = pBuffer;

    while (toWrite != 0) {
        ULONG maxSz = m_ivshmem.chunkSize - m_ulChunkOffset;
        ULONG writeSize = toWrite > maxSz ? maxSz : toWrite;

        // due to something strange, code in qemu reads 1 chunk forward.
        UINT16 chunkIdx = m_ivshmem.chunkIdx + 1;
        if (chunkIdx >= m_ivshmem.maxChunks)
            chunkIdx = 0;
        RtlCopyMemory(&shmem_ptr[chunkIdx * m_ivshmem.chunkSize + m_ulChunkOffset], pReadPtr, writeSize);
        pReadPtr += writeSize;
        toWrite -= writeSize;
        m_ulChunkOffset += writeSize;
        if (m_ulChunkOffset < m_ivshmem.chunkSize) {
            break; // Current chunk is still incomplete
        }

        m_ulChunkOffset = 0;

        KeMemoryBarrier();

        // Send a chunk
        m_ivshmem.chunkIdx = chunkIdx;
        m_pStreamHdr->chunkIdx = m_ivshmem.chunkIdx;
    }
} // WriteData

//=============================================================================
void CIVSHMEMSaveData::ReadData(IN PBYTE pBuffer, IN ULONG ulByteCount)
{
    if (!m_pStreamHdr) {
        return;
    }

    if (!m_pStreamHdr->started) {
        return;
    }

    if (0 == ulByteCount) {
        return;
    }

    if (ulByteCount > (m_ivshmem.bufferSize / 2)) {
        return;
    }

    PBYTE shmem_ptr = (PBYTE)m_pSharedMem + m_pStreamHdr->offset;

    ULONG toRead = ulByteCount;
    PBYTE pReadPtr = pBuffer;

    while (toRead != 0) {
        if (m_ivshmem.chunkIdx == m_pStreamHdr->chunkIdx ||
            m_pStreamHdr->chunkIdx >= m_pStreamHdr->maxChunks) {
            DPF(D_TERSE, ("ReadData: buffer is empty"));
            break;
        }

        KeMemoryBarrier();

        ULONG maxSz = m_ivshmem.chunkSize - m_ulChunkOffset;
        const ULONG readSize = toRead > maxSz ? maxSz : toRead;

        UINT16 chunkIdx = m_ivshmem.chunkIdx;
        RtlCopyMemory(pReadPtr, &shmem_ptr[chunkIdx * m_ivshmem.chunkSize + m_ulChunkOffset], readSize);
        pReadPtr += readSize;
        toRead -= readSize;
        m_ulChunkOffset += readSize;

        if (m_ulChunkOffset < m_ivshmem.chunkSize) {
            break; // Current chunk is still incomplete
        }

        m_ulChunkOffset = 0;

        chunkIdx = m_ivshmem.chunkIdx + 1;
        if (chunkIdx >= m_ivshmem.maxChunks)
            chunkIdx = 0;

        m_ivshmem.chunkIdx = chunkIdx;
    }
}

ULONG CIVSHMEMSaveData::ReadableDataLength(VOID)
{
    if (!m_fCapture)
        return 0;
    return (m_pStreamHdr->chunkIdx + m_pStreamHdr->maxChunks - m_ivshmem.chunkIdx) % m_pStreamHdr->maxChunks * m_ivshmem.chunkSize - m_ulChunkOffset;
}
